module Printer(printExpr, printStmt, printFunction, printProgram) where

import Text.PrettyPrint.HughesPJ
import Types


printExpr :: Expr -> String
printExpr e = render (expr e)


expr :: Expr -> Doc
expr (Ci x)              = integer x
expr (C x)               = double x
expr (V x)               = text x
expr (S x)               = text x
expr (Op "plus" x y)     = prec 4 x <+> char '+' <+> prec 4 y
expr (Op "minus" x y)    = prec 4 x <+> char '-' <+> prec 5 y
expr (Op "mult" x y)     = prec 5 x <+> char '*' <+> prec 5 y
expr (Op "div" x y)      = prec 5 x <+> char '/' <+> prec 6 y
expr (Op f x y)          = text f <> (parens $ commas $ map expr $ [x, y])
expr (Func "lt" [x, y])  = expr x <+> char '<' <+> expr y
expr (Func "gt" [x, y])  = expr x <+> char '>' <+> expr y
expr (Func "eq" [x, y])  = prec 2 x <+> text "==" <+> prec 3 y
expr (Func "and" [x, y]) = prec 1 x <+> text "&&" <+> prec 1 y
expr (Func "or" [x, y])  = prec 0 x <+> text "||" <+> prec 0 y
expr (Func "not" [x])    = char '!' <> prec 10 x
expr (Func f as)         = text f <> (parens $ commas $ map expr $ as)
expr (Cast t e)          = text (typeName t) <> (parens $ expr e)
expr (Dot e x)           = expr e <> char '.' <> text x
expr T                   = text "true"
expr F                   = text "false"

commas :: [Doc] -> Doc
commas es = hsep (punctuate comma es)

semis :: [Doc] -> Doc
semis es = hsep (punctuate semi es)

prec :: Int -> Expr -> Doc
prec n x = if n > p x then parens (expr x) else expr x where
    p (Func "not" _)   = 2
    p (Func "or" _)   = 0
    p (Func "and" _)  = 1

    p (Func "eq" _)   = 2
    p (Func "lt" _)   = 2
    p (Func "gt" _)   = 2

    p (Op "plus"_ _)  = 4
    p (Op "minus"_ _) = 4
    p (Op "mult"_ _)  = 5
    p (Op "div"_ _)   = 5
    
    p _ = 10


printProgram :: Program -> String
printProgram p = render (vcat $ map func p)

    
printFunction :: Function -> String
printFunction f = render (func f)


printStmt :: Stmt -> String
printStmt s = render (stmt s)


func :: Function -> Doc
func (Function f t as b) = text (typeName t) <+> text f <> (parens $ commas $ map arg as) $$ stmt b
    where arg (x, t) = text (typeName t) <+> text x


stmt :: Stmt -> Doc
stmt (Group xs) = char '{' $+$ (nest 4 $ vcat $ map stmt $ xs) $+$ char '}'
stmt (Assign x e) = text x <+> equals <+> expr e
stmt (TypeDecl x t) = text (typeName t) <+> text x
stmt (Return e t) = text "return" <+> expr e <> char ':' <+> text (typeName t)
stmt (Branch c s (Group [])) = text "if" <+> parens (expr c) $$ stmt s
stmt (Branch c s1 s2) = text "if" <+> parens (expr c) $$ stmt s1 $$ text "else"  $$ stmt s2
stmt (While c s) = text "while" <+> parens (expr c) $$ stmt s
stmt (For c u s) = text "for" <+> parens (semis [empty, expr c, stmt u]) $$ stmt s
stmt (FunDef f) = func f
stmt Break = text "break"
stmt Continue = text "continue"


typeName :: Type -> String
typeName Float      = "float"
typeName Float2     = "float2"
typeName Float3     = "float3"
typeName Float4     = "float4"
typeName Boolean    = "bool"
typeName Int        = "int"
typeName Dataset    = "data"
typeName NoType     = "notype"



